import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, input_size, output_size, hidden_size):
        super(Model, self).__init__()
        self.dense1 = nn.Linear(input_size, hidden_size)
        self.dense2 = nn.Linear(hidden_size, output_size)

    def forward(self, obs):
        output = torch.tanh(self.dense1(obs))
        output = torch.softmax(self.dense2(output), dim=-1)
        return output

if __name__ == '__main__':
    model = Model(input_size=2, output_size=4, hidden_size=6)
    out = model.forward(torch.randn((3, 2)))
    print(out)